package storage

import (
	"context"
	"fmt"
	"io"
	"log"
	"log/slog"
	cpb "quipus/proto/gen/common"

	"strings"

	"github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
	"google.golang.org/protobuf/encoding/protojson"
)

type DocVector struct {
	Urn     string    `json:"urn,omitempty"`
	SrcName string    `json:"src_name,omitempty"`
	TgtName string    `json:"tgt_name,omitempty"`
	Content string    `json:"content,omitempty"`
	Vectors []float32 `json:"vectors,omitempty"`
}

type ChunkVector struct {
	Urn     string    `json:"urn,omitempty"`
	RsId    int32     `json:"rs_id,omitempty"`
	Content string    `json:"content,omitempty"`
	Vectors []float32 `json:"vectors,omitempty"`
	Score   float32   `json:"score,omitempty"`
}

type MilvusService struct {
	Client        client.Client
	Ctx           context.Context
	Dim           int
	ContentLength int64
}

type VectorDBConfig struct {
	DBType        string `yaml:"db_type"`
	Url           string `yaml:"url"`
	Username      string `yaml:"username"`
	Password      string `yaml:"password"`
	Dim           int    `yaml:"dim"`
	ContentLength int64  `yaml:"content_length"`
}

type ChunkSummary struct {
	ChunkId string        `json:"chunk_id,omitempty"`
	RsId    int32         `json:"rs_id,omitempty"`
	Summary *cpb.TopicDef `json:"summary,omitempty"`
}

type BuildLog struct {
	ChunkId string `json:"chunk_id,omitempty"`
	RsId    int32  `json:"rs_id,omitempty"`
	Phase   string `json:"phase,omitempty"`
	Status  string `json:"status,omitempty"`
	Reason  string `json:"result,omitempty"`
}

const (
	UrnCol, SrcCol, TgtCol, RidCol, ContentCol, VectorCol, EntityCol = "Urn", "SrcName", "TgtName", "Rid", "Content", "Vector", "Entity"
	TitleCol, SummaryCol                                             = "Title", "Summary"
	PhaseCol                                                         = "Phase"
	ReasonCol                                                        = "Reason"
	StatusCol                                                        = "Status"
)

const (
	EntityPrefix         = "entity_"
	RelationPrefix       = "relation_"
	ChunkPrefix          = "chunk_"
	DocSummaryPrefix     = "doc_summary_"
	GranarySummaryPrefix = "gra_summary_"
	BuildLogPrefix       = "build_"
)

func NewMilvusService(ctx context.Context, config *VectorDBConfig) (*MilvusService, error) {
	c, err := client.NewClient(
		ctx,
		client.Config{
			Address:  config.Url,
			Username: config.Username,
			Password: config.Password,
		},
	)
	if err != nil {
		slog.Error("failed to connect to milvus:", slog.Any("error:", err))
		return nil, err
	}
	return &MilvusService{Client: c, Ctx: ctx, Dim: config.Dim, ContentLength: int64(config.ContentLength)}, nil
}

func (m *MilvusService) DeleteCollection(docId string, prefix string) error {
	collectionName := fmt.Sprintf("%s%s", prefix, docId)
	has, err := m.Client.HasCollection(m.Ctx, collectionName)
	if err != nil {
		slog.Error("failed to check whether collection exists:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
	}
	if has {
		err = m.Client.DropCollection(m.Ctx, collectionName)
		if err != nil {
			slog.Error("failed to drop collection:", slog.Any("error:", err))
		}
	}
	return nil
}

func (m *MilvusService) CreateChunkCollection(docId string, drop bool) error {
	collectionName := fmt.Sprintf("%s%s", ChunkPrefix, docId)
	has, err := m.Client.HasCollection(m.Ctx, collectionName)
	if err != nil {
		slog.Error("failed to check whether collection exists:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
	}
	if has {
		if drop {
			// collection with same name exist, clean up mess

			err = m.Client.DropCollection(m.Ctx, collectionName)
			if err != nil {
				slog.Error("failed to drop collection:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
			}
		} else {
			return nil
		}
	}

	schema := entity.NewSchema().
		WithName(collectionName).
		WithDescription("this is the chunk collection for insert and search").WithAutoID(true).
		WithField(entity.NewField().WithName(UrnCol).WithDataType(entity.FieldTypeVarChar).WithIsPrimaryKey(true).WithMaxLength(70)). // 指定主键的最大长度
		WithField(entity.NewField().WithName(RidCol).WithDataType(entity.FieldTypeInt32)).
		WithField(entity.NewField().WithName(ContentCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(m.ContentLength)). // 为 SrcCol 指定最大长度
		WithField(entity.NewField().WithName(VectorCol).WithDataType(entity.FieldTypeFloatVector).WithDim(int64(m.Dim)))

	err = m.Client.CreateCollection(m.Ctx, schema, 1) // only 1 shard
	if err != nil {
		slog.Error("failed to create collection:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
		return err
	}

	m.CreateIndex(collectionName)
	return nil
}

func (m *MilvusService) DropCollection(collectionName string) error {
	slog.Info("drop collection:", collectionName)

	err := m.Client.DropCollection(m.Ctx, collectionName)
	if err != nil {
		slog.Error("failed to drop collection:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
	}
	return err
}

func (m *MilvusService) CreateEntityCollection(docId string, drop bool) error {
	collectionName := fmt.Sprintf("%s%s", EntityPrefix, docId)

	has, err := m.Client.HasCollection(m.Ctx, collectionName)
	if err != nil {
		slog.Error("failed to check whether collection exists:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
		return nil
	}
	if has {
		if drop {
			// collection with same name exist, clean up mess
			err := m.DropCollection(collectionName)
			if err != nil {
				return err
			}
		} else {
			return nil
		}
	}

	schema := entity.NewSchema().
		WithName(collectionName).
		WithDescription("this is the entity collection for insert and search").
		WithField(entity.NewField().WithName(UrnCol).WithDataType(entity.FieldTypeVarChar).WithIsPrimaryKey(true).WithMaxLength(256)).WithAutoID(false). // 指定主键的最大长度
		WithField(entity.NewField().WithName(EntityCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(100)).                                       // 为 SrcCol 指定最大长度
		WithField(entity.NewField().WithName(ContentCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(m.ContentLength)).                          // 为 ContentCol 指定最大长度
		WithField(entity.NewField().WithName(VectorCol).WithDataType(entity.FieldTypeFloatVector).WithDim(int64(m.Dim)))

	err = m.Client.CreateCollection(m.Ctx, schema, 1) // only 1 shard
	if err != nil {
		slog.Error("failed to create collection:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
		return err
	}

	m.CreateIndex(collectionName)

	slog.Info("create collection success", slog.Any("collectionName:", collectionName))

	return nil
}

func (m *MilvusService) CreateRelationCollection(docId string, drop bool) error {
	collectionName := fmt.Sprintf("%s%s", RelationPrefix, docId)

	has, err := m.Client.HasCollection(m.Ctx, collectionName)
	if err != nil {
		slog.Error("failed to check whether collection exists:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
		return nil
	}
	if has {
		if drop {
			// collection with same name exist, clean up mess
			err := m.DropCollection(collectionName)
			if err != nil {
				return err
			}
		} else {
			return nil
		}
	}

	schema := entity.NewSchema().
		WithName(collectionName).
		WithDescription("this is the relation collection for insert and search").
		WithField(entity.NewField().WithName(UrnCol).WithDataType(entity.FieldTypeVarChar).WithIsPrimaryKey(true).WithMaxLength(256)). // 指定主键的最大长度
		WithField(entity.NewField().WithName(SrcCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(100)).                        // 为 SrcCol 指定最大长度
		WithField(entity.NewField().WithName(TgtCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(100)).                        // 为 TgtCol 指定最大长度
		WithField(entity.NewField().WithName(ContentCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(m.ContentLength)).        // 为 ContentCol 指定最大长度
		WithField(entity.NewField().WithName(VectorCol).WithDataType(entity.FieldTypeFloatVector).WithDim(int64(m.Dim)))

	err = m.Client.CreateCollection(m.Ctx, schema, 1) // only 1 shard
	if err != nil {
		slog.Error("failed to create collection:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
		return err
	}

	m.CreateIndex(collectionName)

	slog.Info("create collection success", slog.Any("collectionName:", collectionName))

	return nil
}

func (m *MilvusService) CreateDocSummaryCollection(docId string, drop bool) error {
	collectionName := fmt.Sprintf("%s%s", DocSummaryPrefix, docId)

	has, err := m.Client.HasCollection(m.Ctx, collectionName)
	if err != nil {
		slog.Error("failed to check whether collection exists:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
		return nil
	}
	if has {
		if drop {
			// collection with same name exist, clean up mess

			err := m.DropCollection(collectionName)
			if err != nil {
				return err
			}
		} else {
			return nil
		}
	}

	schema := entity.NewSchema().
		WithName(collectionName).
		WithDescription("this is the doc summary collection for insert and search").
		WithField(entity.NewField().WithName(UrnCol).WithDataType(entity.FieldTypeVarChar).WithIsPrimaryKey(true).WithMaxLength(256)). // 指定主键的最大长度
		WithField(entity.NewField().WithName(RidCol).WithDataType(entity.FieldTypeInt32)).                                             // 为 SrcCol 指定最大长度
		WithField(entity.NewField().WithName(TitleCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(200)).                      // 为 SrcCol 指定最大长度
		WithField(entity.NewField().WithName(SummaryCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(m.ContentLength)).        // 为 TgtCol 指定最大长度
		WithField(entity.NewField().WithName(ContentCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(m.ContentLength)).        // 为 ContentCol 指定最大长度
		WithField(entity.NewField().WithName(VectorCol).WithDataType(entity.FieldTypeFloatVector).WithDim(int64(m.Dim)))

	err = m.Client.CreateCollection(m.Ctx, schema, 1) // only 1 shard
	if err != nil {
		slog.Error("failed to create collection:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
		return err
	}

	m.CreateIndex(collectionName)

	return nil
}

func (m *MilvusService) CreateGSummaryCollection(docId string, drop bool) error {
	collectionName := fmt.Sprintf("%s%s", GranarySummaryPrefix, docId)

	has, err := m.Client.HasCollection(m.Ctx, collectionName)
	if err != nil {
		slog.Error("failed to check whether collection exists:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
		return nil
	}
	if has {
		if drop {
			// collection with same name exist, clean up mess

			err := m.DropCollection(collectionName)
			if err != nil {
				return err
			}
		} else {
			return nil
		}
	}

	schema := entity.NewSchema().
		WithName(collectionName).
		WithDescription("this is the granary summary collection for insert and search").
		WithField(entity.NewField().WithName(UrnCol).WithDataType(entity.FieldTypeVarChar).WithIsPrimaryKey(true).WithMaxLength(256)). // 指定主键的最大长度
		WithField(entity.NewField().WithName(TitleCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(200)).                      // 为 SrcCol 指定最大长度
		WithField(entity.NewField().WithName(SummaryCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(m.ContentLength)).        // 为 TgtCol 指定最大长度
		WithField(entity.NewField().WithName(ContentCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(m.ContentLength)).        // 为 ContentCol 指定最大长度
		WithField(entity.NewField().WithName(VectorCol).WithDataType(entity.FieldTypeFloatVector).WithDim(int64(m.Dim)))

	err = m.Client.CreateCollection(m.Ctx, schema, 1) // only 1 shard
	if err != nil {
		slog.Error("failed to create collection:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
		return err
	}

	m.CreateIndex(collectionName)
	return nil
}

func (m *MilvusService) CreateIndex(collectionName string) {
	// idx, err := entity.NewIndexIvfFlat(entity.L2, 2)
	idx, err := entity.NewIndexFlat(entity.L2)
	if err != nil {
		slog.Error("fail to create ivf flat index:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
	}

	err = m.Client.CreateIndex(m.Ctx, collectionName, VectorCol, idx, false)
	if err != nil {
		slog.Error("fail to create index:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
	}
}

func (m *MilvusService) InsertRelation(docId string, entities []*DocVector, embed_func func(context.Context, string) []float32) (int, error) {
	collectionName := fmt.Sprintf("%s%s", RelationPrefix, docId)
	urns := make([]string, 0, len(entities))
	srcs := make([]string, 0, len(entities))
	tgts := make([]string, 0, len(entities))
	contents := make([]string, 0, len(entities))
	vectors := make([][]float32, 0, len(entities))
	for _, entity := range entities {
		urns = append(urns, entity.Urn)
		srcs = append(srcs, entity.SrcName)
		tgts = append(tgts, entity.TgtName)
		contents = append(contents, entity.Content)
		v := embed_func(m.Ctx, entity.Content)
		entity.Vectors = v
		// slog.Info("doc info:", slog.Any("entity:", entity))

		vectors = append(vectors, v)
	}

	urnColumn := entity.NewColumnVarChar(UrnCol, urns)
	srcColumn := entity.NewColumnVarChar(SrcCol, srcs)
	tgtColumn := entity.NewColumnVarChar(TgtCol, tgts)
	contentColumn := entity.NewColumnVarChar(ContentCol, contents)
	vectorColumn := entity.NewColumnFloatVector(VectorCol, m.Dim, vectors)
	e, err := m.Client.Insert(m.Ctx, collectionName, "", urnColumn, srcColumn, tgtColumn, contentColumn, vectorColumn)
	if err != nil {
		slog.Error("failed to insert relation data:", slog.Any("error:", err))
		return 0, err
	}

	if err := m.Client.Flush(m.Ctx, collectionName, false); err != nil {
		slog.Error("failed to flush data, ", slog.Any("error:", err))
	}

	m.CreateIndex(collectionName)

	return e.Len(), nil
}

func (m *MilvusService) InsertEntity(docId string, entities []*DocVector, embed_func func(context.Context, string) []float32) (int, error) {
	collectionName := fmt.Sprintf("%s%s", EntityPrefix, docId)

	urns := make([]string, 0, len(entities))
	srcs := make([]string, 0, len(entities))
	contents := make([]string, 0, len(entities))
	vectors := make([][]float32, 0, len(entities))
	for _, entity := range entities {
		urns = append(urns, entity.Urn)
		srcs = append(srcs, entity.SrcName)
		contents = append(contents, entity.Content)
		v := embed_func(m.Ctx, entity.Content)
		entity.Vectors = v
		vectors = append(vectors, v)
	}

	urnColumn := entity.NewColumnVarChar(UrnCol, urns)
	srcColumn := entity.NewColumnVarChar(EntityCol, srcs)
	contentColumn := entity.NewColumnVarChar(ContentCol, contents)
	vectorColumn := entity.NewColumnFloatVector(VectorCol, m.Dim, vectors)
	e, err := m.Client.Insert(m.Ctx, collectionName, "", urnColumn, srcColumn, contentColumn, vectorColumn)
	if err != nil {
		slog.Error("failed to insert entity data:", slog.Any("entity:", err))
		return 0, err
	}

	if err := m.Client.Flush(m.Ctx, collectionName, false); err != nil {
		slog.Error("failed to flush data:", slog.Any("error:", err))
	}
	m.CreateIndex(collectionName)
	return e.Len(), nil
}

func (m *MilvusService) InsertChunk(docId string, chunks []*ChunkVector, embed_func func(context.Context, string) []float32) (int, error) {
	collectionName := fmt.Sprintf("%s%s", ChunkPrefix, docId)
	urns := make([]string, 0, len(chunks))
	contents := make([]string, 0, len(chunks))
	vectors := make([][]float32, 0, len(chunks))
	rsIds := make([]int32, 0, len(chunks))
	for _, chunk := range chunks {
		urns = append(urns, chunk.Urn)
		contents = append(contents, chunk.Content)
		v := embed_func(m.Ctx, chunk.Content)
		if v == nil {
			slog.Error("failed to embed content:", slog.Any("content:", chunk.Content))
			return 0, fmt.Errorf("failed to embed content: %s", chunk.Content)
		}
		chunk.Vectors = v

		vectors = append(vectors, v)
		rsIds = append(rsIds, chunk.RsId)
	}

	urnColumn := entity.NewColumnVarChar(UrnCol, urns)
	contentColumn := entity.NewColumnVarChar(ContentCol, contents)
	ridColumn := entity.NewColumnInt32(RidCol, rsIds)

	vectorColumn := entity.NewColumnFloatVector(VectorCol, m.Dim, vectors)
	e, err := m.Client.Insert(m.Ctx, collectionName, "", urnColumn, ridColumn, contentColumn, vectorColumn)
	if err != nil {
		slog.Info("content:", slog.Any("content:", contents))
		slog.Error("failed to insert chunk data:", slog.Any("error:", err))
		return 0, err
	}

	if err := m.Client.Flush(m.Ctx, collectionName, false); err != nil {
		slog.Error("failed to flush data:", slog.Any("error:", err))
	}
	m.CreateIndex(collectionName)
	return e.Len(), nil
}

func (m *MilvusService) GetDocSummary(docId string, rsId int32, ids []string) ([]*ChunkSummary, error) {
	collectionName := fmt.Sprintf("%s%s", DocSummaryPrefix, docId)

	err := m.LoadCollection(collectionName)
	if err != nil {
		if strings.Contains(err.Error(), "index not found") {
			return nil, nil
		}
		return nil, err
	}

	urnColumn := entity.NewColumnVarChar(UrnCol, ids)

	results, err := m.Client.Get(m.Ctx, collectionName, urnColumn)
	if err != nil {
		return nil, err
	}

	segments := []*ChunkSummary{}

	for _, field := range results {

		switch field.Name() {
		case UrnCol:
			c, ok := field.(*entity.ColumnVarChar)
			if !ok {
				continue
			}
			for i := 0; i < c.Len(); i++ {
				v, e := c.ValueByIdx(i)
				if e != nil {
					continue
				}

				if len(segments) < i+1 {
					segments = append(segments, &ChunkSummary{ChunkId: v})
				} else {
					segments[i].ChunkId = v
				}
			}
		case ContentCol:
			c, ok := field.(*entity.ColumnVarChar)
			if !ok {
				continue
			}
			for i := 0; i < c.Len(); i++ {
				v, e := c.ValueByIdx(i)
				if e != nil {
					continue
				}
				summary := &cpb.TopicDef{}
				protojson.Unmarshal([]byte(v), summary)
				if len(segments) < i+1 {

					chunk := &ChunkSummary{}
					chunk.Summary = summary
					segments = append(segments, chunk)
				} else {
					segments[i].Summary = summary
				}
			}
		default:
			slog.Warn("ignore field:", slog.Any("field:", field.Name()))

		}
	}
	return segments, nil
}

func (m *MilvusService) InsertDocSummary(docId string, chunks []*ChunkSummary, embed_func func(context.Context, string) []float32) (int, error) {
	collectionName := fmt.Sprintf("%s%s", DocSummaryPrefix, docId)
	urns := make([]string, 0, len(chunks))
	titiles := make([]string, 0, len(chunks))
	summaries := make([]string, 0, len(chunks))
	contents := make([]string, 0, len(chunks))
	rsIds := make([]int32, 0, len(chunks))

	vectors := make([][]float32, 0, len(chunks))
	for _, chunk := range chunks {
		urns = append(urns, chunk.ChunkId)
		titiles = append(titiles, chunk.Summary.Primary.Title)
		summaries = append(summaries, chunk.Summary.Primary.Summary)
		content, e := protojson.Marshal(chunk.Summary)
		// secondSummary, e := protojson.Marshal(chunk.Summary.Units)
		if e != nil {
			slog.Error("failed to marshal second summary:", slog.Any("error:", e))
			return 0, e
		}
		contents = append(contents, string(content))
		v := embed_func(m.Ctx, string(content))
		vectors = append(vectors, v)
		rsIds = append(rsIds, chunk.RsId)
	}

	urnColumn := entity.NewColumnVarChar(UrnCol, urns)
	titleColumn := entity.NewColumnVarChar(TitleCol, titiles)
	summaryColumn := entity.NewColumnVarChar(SummaryCol, summaries)
	contentColumn := entity.NewColumnVarChar(ContentCol, contents)
	vectorColumn := entity.NewColumnFloatVector(VectorCol, m.Dim, vectors)
	rsIdColumn := entity.NewColumnInt32(RidCol, rsIds)

	e, err := m.Client.Insert(m.Ctx, collectionName, "", urnColumn, rsIdColumn, titleColumn, summaryColumn, contentColumn, vectorColumn)
	if err != nil {
		slog.Error("failed to insert doc summary data:", slog.Any("error:", err))
		return 0, err
	}

	if err := m.Client.Flush(m.Ctx, collectionName, false); err != nil {
		slog.Error("failed to flush data:", slog.Any("error:", err))
	}
	m.CreateIndex(collectionName)
	return e.Len(), nil
}

func (m *MilvusService) InsertGSummary(docId string, chunks []*ChunkSummary, embed_func func(context.Context, string) []float32) (int, error) {
	collectionName := fmt.Sprintf("%s%s", GranarySummaryPrefix, docId)
	urns := make([]string, 0, len(chunks))
	titiles := make([]string, 0, len(chunks))
	summaries := make([]string, 0, len(chunks))
	contents := make([]string, 0, len(chunks))

	vectors := make([][]float32, 0, len(chunks))
	for _, chunk := range chunks {
		urns = append(urns, chunk.ChunkId)
		titiles = append(titiles, chunk.Summary.Primary.Title)
		summaries = append(summaries, chunk.Summary.Primary.Summary)
		content, e := protojson.Marshal(chunk.Summary)
		// secondSummary, e := protojson.Marshal(chunk.Summary.Units)
		if e != nil {
			slog.Error("failed to marshal second summary:", slog.Any("error:", e))
			return 0, e
		}
		contents = append(contents, string(content))
		v := embed_func(m.Ctx, string(content))
		vectors = append(vectors, v)
	}

	urnColumn := entity.NewColumnVarChar(UrnCol, urns)
	titleColumn := entity.NewColumnVarChar(TitleCol, titiles)
	summaryColumn := entity.NewColumnVarChar(SummaryCol, summaries)
	contentColumn := entity.NewColumnVarChar(ContentCol, contents)
	vectorColumn := entity.NewColumnFloatVector(VectorCol, m.Dim, vectors)

	e, err := m.Client.Insert(m.Ctx, collectionName, "", urnColumn, titleColumn, summaryColumn, contentColumn, vectorColumn)
	if err != nil {
		slog.Error("failed to insert granary summary data:", slog.Any("error:", err))
		return 0, err
	}

	if err := m.Client.Flush(m.Ctx, collectionName, false); err != nil {
		slog.Error("failed to flush data:", slog.Any("error:", err))
	}
	m.CreateIndex(collectionName)
	return e.Len(), nil
}

func (m *MilvusService) LoadCollection(collectionName string) error {
	err := m.Client.LoadCollection(m.Ctx, collectionName, false)
	if err != nil {
		slog.Error("failed to load collection:", slog.Any("error:", err))
		return err
	}
	return nil
}

func (m *MilvusService) ReleaseCollection(collectionName string) error {
	err := m.Client.ReleaseCollection(m.Ctx, collectionName)
	if err != nil {
		slog.Error("failed to release collection:", slog.Any("error:", err))

		return err
	}
	return nil
}

func (m *MilvusService) QueryRelation(query string, docId string, embed_func func(context.Context, string) []float32) ([]*DocVector, error) {
	collectionName := fmt.Sprintf("%s%s", RelationPrefix, docId)

	qv := embed_func(m.Ctx, query)

	vector := entity.FloatVector(qv)
	err := m.LoadCollection(collectionName)
	if err != nil {
		return nil, err
	}

	// Use flat search param
	sp, _ := entity.NewIndexFlatSearchParam()
	// sp, _ := entity.NewIndexIvfFlatSearchParam(16)
	sr, err := m.Client.Search(m.Ctx, collectionName, []string{}, "", []string{UrnCol, SrcCol, TgtCol, ContentCol}, []entity.Vector{vector}, VectorCol,
		entity.L2, 10, sp)
	if err != nil {
		slog.Error("fail to search collection:", slog.Any("error:", err))
		return nil, err
	}

	docs := []*DocVector{}

	for _, result := range sr {
		segments := []*DocVector{}
		for _, field := range result.Fields {
			switch field.Name() {
			case UrnCol:
				c, ok := field.(*entity.ColumnVarChar)
				if !ok {
					continue
				}
				for i := 0; i < c.Len(); i++ {
					v, e := c.ValueByIdx(i)
					if e != nil {
						continue
					}

					if len(segments) < i+1 {
						segments = append(segments, &DocVector{Urn: v})
					} else {
						segments[i].Urn = v
					}
				}
			case SrcCol:
				c, ok := field.(*entity.ColumnVarChar)
				if !ok {
					continue
				}
				for i := 0; i < c.Len(); i++ {
					v, e := c.ValueByIdx(i)
					if e != nil {
						continue
					}

					if len(segments) < i+1 {
						segments = append(segments, &DocVector{SrcName: v})
					} else {
						segments[i].SrcName = v
					}
				}
			case TgtCol:
				c, ok := field.(*entity.ColumnVarChar)
				if !ok {
					continue
				}
				for i := 0; i < c.Len(); i++ {
					v, e := c.ValueByIdx(i)
					if e != nil {
						continue
					}

					if len(segments) < i+1 {
						segments = append(segments, &DocVector{TgtName: v})
					} else {
						segments[i].TgtName = v
					}
				}
			case ContentCol:
				c, ok := field.(*entity.ColumnVarChar)
				if !ok {
					continue
				}
				for i := 0; i < c.Len(); i++ {
					v, e := c.ValueByIdx(i)
					if e != nil {
						continue
					}

					if len(segments) < i+1 {
						segments = append(segments, &DocVector{Content: v})
					} else {
						segments[i].Content = v
					}
				}
			default:
				log.Println("QueryRelation unknow field:", slog.Any("field:", field.Name()))
			}
		}

		docs = append(docs, segments...)
	}

	return docs, nil
}

func (m *MilvusService) QueryEntity(query string, docId string, embed_func func(context.Context, string) []float32) ([]*DocVector, error) {
	collectionName := fmt.Sprintf("%s%s", EntityPrefix, docId)

	qv := embed_func(m.Ctx, query)

	vector := entity.FloatVector(qv)
	err := m.LoadCollection(collectionName)
	if err != nil {
		return nil, err
	}
	// Use flat search param
	sp, _ := entity.NewIndexFlatSearchParam()
	// sp, _ := entity.NewIndexIvfFlatSearchParam(16)
	sr, err := m.Client.Search(m.Ctx, collectionName, []string{}, "", []string{UrnCol, EntityCol, ContentCol}, []entity.Vector{vector}, VectorCol,
		entity.L2, 10, sp)
	if err != nil {
		slog.Error("fail to search collection:", slog.Any("error:", err))

		return nil, err
	}

	docs := []*DocVector{}

	for _, result := range sr {
		segments := []*DocVector{}
		for _, field := range result.Fields {
			switch field.Name() {
			case UrnCol:
				c, ok := field.(*entity.ColumnVarChar)
				if !ok {
					continue
				}
				for i := 0; i < c.Len(); i++ {
					v, e := c.ValueByIdx(i)
					if e != nil {
						continue
					}

					if len(segments) < i+1 {
						segments = append(segments, &DocVector{Urn: v})
					} else {
						segments[i].Urn = v
					}
				}
			case EntityCol:
				c, ok := field.(*entity.ColumnVarChar)
				if !ok {
					continue
				}
				for i := 0; i < c.Len(); i++ {
					v, e := c.ValueByIdx(i)
					if e != nil {
						continue
					}

					if len(segments) < i+1 {
						segments = append(segments, &DocVector{SrcName: v})
					} else {
						segments[i].SrcName = v
					}
				}
			case ContentCol:
				c, ok := field.(*entity.ColumnVarChar)
				if !ok {
					continue
				}
				for i := 0; i < c.Len(); i++ {
					v, e := c.ValueByIdx(i)
					if e != nil {
						continue
					}

					if len(segments) < i+1 {
						segments = append(segments, &DocVector{Content: v})
					} else {
						segments[i].Content = v
					}
				}
			default:
				slog.Error("QueryEntity unknow field:", slog.Any("field:", field.Name()))
			}
		}

		docs = append(docs, segments...)
	}

	return docs, nil
}

type SearchResult struct {
	CollectionName string
	Score          []float32
	Data           client.ResultSet
}

func (m *MilvusService) QueryResult(query string, docId string, queryCol []string, embed_func func(context.Context, string) []float32) ([]*SearchResult, error) {
	collectionName := fmt.Sprintf("%s%s", ChunkPrefix, docId)

	qv := embed_func(m.Ctx, query)

	vector := entity.FloatVector(qv)

	// Use flat search param
	sp, _ := entity.NewIndexFlatSearchParam()
	// sp, _ := entity.NewIndexIvfFlatSearchParam(16)
	sr, err := m.Client.Search(m.Ctx, collectionName, []string{}, "", queryCol, []entity.Vector{vector}, VectorCol,
		entity.L2, 5, sp)
	if err != nil {
		slog.Error("fail to search collection:", slog.Any("error:", err))

		return nil, err
	}

	var results []*SearchResult
	for _, hit := range sr {
		results = append(results, &SearchResult{
			CollectionName: collectionName,
			Score:          hit.Scores,
			Data:           hit.Fields, // 可附加业务数据
		})
	}

	return results, nil
}

func (m *MilvusService) QueryChunk(query string, docId string, embed_func func(context.Context, string) []float32) ([]*ChunkVector, error) {
	collectionName := fmt.Sprintf("%s%s", ChunkPrefix, docId)

	qv := embed_func(m.Ctx, query)

	vector := entity.FloatVector(qv)

	// Use flat search param
	sp, _ := entity.NewIndexFlatSearchParam()
	// sp, _ := entity.NewIndexIvfFlatSearchParam(16)
	sr, err := m.Client.Search(m.Ctx, collectionName, []string{}, "", []string{UrnCol, ContentCol}, []entity.Vector{vector}, VectorCol,
		entity.L2, 5, sp)
	if err != nil {
		slog.Error("fail to search collection:", slog.Any("error:", err))

		return nil, err
	}

	chunks := []*ChunkVector{}

	for _, result := range sr {
		segments := []*ChunkVector{}

		for _, field := range result.Fields {
			switch field.Name() {
			case UrnCol:
				c, ok := field.(*entity.ColumnVarChar)
				if !ok {
					continue
				}
				for i := 0; i < c.Len(); i++ {
					v, e := c.ValueByIdx(i)
					if e != nil {
						continue
					}

					if len(segments) < i+1 {
						segments = append(segments, &ChunkVector{Urn: v, Score: result.Scores[i]})
					} else {
						segments[i].Urn = v
						segments[i].Score = result.Scores[i]
					}
				}
			case ContentCol:
				c, ok := field.(*entity.ColumnVarChar)
				if !ok {
					continue
				}
				for i := 0; i < c.Len(); i++ {
					v, e := c.ValueByIdx(i)
					if e != nil {
						continue
					}

					if len(segments) < i+1 {
						segments = append(segments, &ChunkVector{Content: v})
					} else {
						segments[i].Content = v
					}
				}
			default:
				log.Println("QueryChunk unknow field:", slog.Any("field:", field.Name()))
			}
		}

		if len(segments) > 0 {
			chunks = append(chunks, segments...)
		}
	}

	return chunks, nil
}

func (m *MilvusService) GetChunk(docId string, ids []string) ([]*ChunkVector, error) {
	collectionName := fmt.Sprintf("%s%s", ChunkPrefix, docId)

	err := m.LoadCollection(collectionName)
	if err != nil {
		if strings.Contains(err.Error(), "index not found") {
			return nil, nil
		}
		return nil, err
	}

	urnColumn := entity.NewColumnVarChar(UrnCol, ids)

	results, err := m.Client.Get(m.Ctx, collectionName, urnColumn)
	if err != nil {
		return nil, err
	}

	chunks := []*ChunkVector{}

	for _, field := range results {

		switch field.Name() {
		case UrnCol:
			c, ok := field.(*entity.ColumnVarChar)
			if !ok {
				continue
			}
			for i := 0; i < c.Len(); i++ {
				v, e := c.ValueByIdx(i)
				if e != nil {
					continue
				}

				if len(chunks) < i+1 {
					chunks = append(chunks, &ChunkVector{Urn: v})
				} else {
					chunks[i].Urn = v
				}
			}
		case ContentCol:
			c, ok := field.(*entity.ColumnVarChar)
			if !ok {
				continue
			}
			for i := 0; i < c.Len(); i++ {
				v, e := c.ValueByIdx(i)
				if e != nil {
					continue
				}

				if len(chunks) < i+1 {
					chunks = append(chunks, &ChunkVector{Content: v})
				} else {
					chunks[i].Content = v
				}
			}
		case VectorCol:
			c, ok := field.(*entity.ColumnFloatVector)
			if !ok {
				continue
			}
			for i := 0; i < c.Len(); i++ {
				v := c.Data()
				if len(v) < 1 {
					continue
				}

				if len(chunks) < i+1 {
					chunks = append(chunks, &ChunkVector{Vectors: v[0]})
				} else {
					chunks[i].Vectors = v[0]
				}
			}
		default:
			slog.Error("GetChunk unknow field:", slog.Any("field:", field.Name()))

		}
	}
	return chunks, nil
}

func (m *MilvusService) QueryChunkIterator(docId string, rsId int32, extractChunk func(docId string, rsId int32, ckv *ChunkVector) error) error {
	collectionName := fmt.Sprintf("%s%s", ChunkPrefix, docId)
	err := m.LoadCollection(collectionName)
	if err != nil {
		if strings.Contains(err.Error(), "index not found") {
			return nil
		}
		return err
	}

	opts := client.NewQueryIteratorOption(collectionName).WithOutputFields(UrnCol, RidCol, ContentCol).WithBatchSize(100)

	if rsId != 0 {
		opts = opts.WithExpr(fmt.Sprintf("%s == %d", RidCol, rsId))
	}

	itr, err := m.Client.QueryIterator(m.Ctx, opts)

	if err != nil {
		return err
	}
queryLoop:
	for {
		select {
		case <-m.Ctx.Done():
			slog.Info("process is stopped")
			return nil
		default:
			{
				rs, err := itr.Next(m.Ctx)
				if err != nil {
					if err == io.EOF {
						log.Println("QueryChunkIterator iterator reach EOF")
						break queryLoop
					}
					slog.Info("failed to query iterator. next: ", slog.Any("err:", err))
					return err
				}

				segments := []*ChunkVector{}

				for _, field := range rs {

					switch field.Name() {
					case UrnCol:
						c, ok := field.(*entity.ColumnVarChar)
						if !ok {
							continue
						}
						for i := 0; i < c.Len(); i++ {
							v, e := c.ValueByIdx(i)
							if e != nil {
								continue
							}

							if len(segments) < i+1 {
								segments = append(segments, &ChunkVector{Urn: v})
							} else {
								segments[i].Urn = v
							}
						}
					case RidCol:
						c, ok := field.(*entity.ColumnInt32)
						if !ok {
							continue
						}
						for i := 0; i < c.Len(); i++ {
							v, e := c.ValueByIdx(i)
							if e != nil {
								continue
							}

							if len(segments) < i+1 {
								segments = append(segments, &ChunkVector{RsId: v})
							} else {
								segments[i].RsId = v
							}
						}
					case ContentCol:
						c, ok := field.(*entity.ColumnVarChar)
						if !ok {
							continue
						}
						for i := 0; i < c.Len(); i++ {
							v, e := c.ValueByIdx(i)
							if e != nil {
								continue
							}

							if len(segments) < i+1 {
								segments = append(segments, &ChunkVector{Content: v})
							} else {
								segments[i].Content = v
							}
						}
					case VectorCol:
						c, ok := field.(*entity.ColumnFloatVector)
						if !ok {
							continue
						}
						for i := 0; i < c.Len(); i++ {
							v := c.Data()
							if len(v) < 1 {
								continue
							}

							if len(segments) < i+1 {
								segments = append(segments, &ChunkVector{Vectors: v[0]})
							} else {
								segments[i].Vectors = v[0]
							}
						}
					default:
						slog.Error("QueryChunkIterator unknow field:", slog.Any("field:", field.Name()))
					}
				}

				if len(segments) > 0 {
					for _, segment := range segments {
						// err := extractChunk(docId, segment.Urn, segment.Content)
						err := extractChunk(docId, rsId, segment)
						if err != nil {
							return err
						}
					}
				}
			}
		}

	}
	return nil
}

func (m *MilvusService) DeleteDocSummary(docId string, rsId int32) error {
	collectionName := fmt.Sprintf("%s%s", DocSummaryPrefix, docId)
	err := m.LoadCollection(collectionName)
	if err != nil {
		if strings.Contains(err.Error(), "index not found") {
			return nil
		}
		return err
	}
	return m.Client.Delete(m.Ctx, collectionName, "", fmt.Sprintf("%s == %d", RidCol, rsId))
}

func (m *MilvusService) QueryDocSummary(docId string, rsId int32) ([]*ChunkSummary, error) {
	collectionName := fmt.Sprintf("%s%s", DocSummaryPrefix, docId)
	err := m.LoadCollection(collectionName)
	if err != nil {
		if strings.Contains(err.Error(), "index not found") {
			return nil, nil
		}
		return nil, err
	}

	opts := client.NewQueryIteratorOption(collectionName).WithOutputFields(UrnCol, ContentCol).WithBatchSize(100)
	if rsId != 0 {
		opts = opts.WithExpr(fmt.Sprintf("%s == %d", RidCol, rsId))
	}

	return m.doQuerySummary(opts)
}

func (m *MilvusService) doQuerySummary(opts *client.QueryIteratorOption) ([]*ChunkSummary, error) {
	itr, err := m.Client.QueryIterator(m.Ctx, opts)

	if err != nil {
		return nil, err
	}

	segments := []*ChunkSummary{}

	for {
		select {
		case <-m.Ctx.Done():
			slog.Info("process is stopped")
			return nil, nil
		default:
			{
				rs, err := itr.Next(m.Ctx)
				if err != nil {
					if err == io.EOF {
						log.Println("QuerySummary iterator reach EOF")
						return segments, nil
					}
					slog.Info("failed to query iterator. next: ", slog.Any("err:", err))
				}

				for _, field := range rs {

					switch field.Name() {
					case UrnCol:
						c, ok := field.(*entity.ColumnVarChar)
						if !ok {
							continue
						}
						for i := 0; i < c.Len(); i++ {
							v, e := c.ValueByIdx(i)
							if e != nil {
								continue
							}

							if len(segments) < i+1 {
								segments = append(segments, &ChunkSummary{ChunkId: v})
							} else {
								segments[i].ChunkId = v
							}
						}
					case ContentCol:
						c, ok := field.(*entity.ColumnVarChar)
						if !ok {
							continue
						}
						for i := 0; i < c.Len(); i++ {
							v, e := c.ValueByIdx(i)
							if e != nil {
								continue
							}
							summary := &cpb.TopicDef{}
							e = protojson.Unmarshal([]byte(v), summary)
							if e != nil {
								continue
							}

							if len(segments) < i+1 {

								chunk := &ChunkSummary{}
								chunk.Summary = summary
								segments = append(segments, chunk)
							} else {
								segments[i].Summary = summary
							}
						}
					default:
						slog.Error("doQuerySummary unknow field:", slog.Any("field:", field.Name()))
					}
				}
			}
		}
	}
}

func (m *MilvusService) QueryGSummary(docId string) ([]*ChunkSummary, error) {
	collectionName := fmt.Sprintf("%s%s", GranarySummaryPrefix, docId)
	err := m.LoadCollection(collectionName)
	if err != nil {
		if strings.Contains(err.Error(), "index not found") {
			return nil, nil
		}
		return nil, err
	}

	opts := client.NewQueryIteratorOption(collectionName).WithOutputFields(UrnCol, ContentCol).WithBatchSize(100)

	return m.doQuerySummary(opts)
}

func (m *MilvusService) CreateBuildLogCollection(docId string, drop bool) error {
	collectionName := fmt.Sprintf("%s%s", BuildLogPrefix, docId)

	has, err := m.Client.HasCollection(m.Ctx, collectionName)
	if err != nil {
		slog.Error("failed to check whether collection exists:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
		return nil
	}
	if has {
		if drop {
			// collection with same name exist, clean up mess

			err := m.DropCollection(collectionName)
			if err != nil {
				return err
			}
		} else {
			return nil
		}
	}

	schema := entity.NewSchema().
		WithName(collectionName).
		WithDescription("this is the build log collection for insert and search").
		WithField(entity.NewField().WithName(UrnCol).WithDataType(entity.FieldTypeVarChar).WithIsPrimaryKey(true).WithMaxLength(256)). // 指定主键的最大长度
		WithField(entity.NewField().WithName(RidCol).WithDataType(entity.FieldTypeInt32)).                                             // 为 SrcCol 指定最大长度
		WithField(entity.NewField().WithName(PhaseCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(200)).
		WithField(entity.NewField().WithName(StatusCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(200)). // 为 SrcCol 指定最大长度
		WithField(entity.NewField().WithName(ReasonCol).WithDataType(entity.FieldTypeVarChar).WithMaxLength(200)). // 为 TgtCol 指定最大长度
		WithField(entity.NewField().WithName(VectorCol).WithDataType(entity.FieldTypeFloatVector).WithDim(int64(m.Dim)))

	err = m.Client.CreateCollection(m.Ctx, schema, 1) // only 1 shard
	if err != nil {
		slog.Error("failed to create collection:", slog.Any("error:", err), slog.Any("collectionName:", collectionName))
		return err
	}

	m.CreateIndex(collectionName)

	return nil
}

func (m *MilvusService) UpsertBuildLog(docId string, blgs []*BuildLog) (int, error) {
	collectionName := fmt.Sprintf("%s%s", BuildLogPrefix, docId)
	urns := make([]string, 0, len(blgs))
	phases := make([]string, 0, len(blgs))
	statuses := make([]string, 0, len(blgs))
	reasons := make([]string, 0, len(blgs))
	rsIds := make([]int32, 0, len(blgs))

	vectors := make([][]float32, 0, len(blgs))
	for _, b := range blgs {
		urns = append(urns, b.ChunkId)
		phases = append(phases, b.Phase)
		reasons = append(reasons, b.Reason)
		v := make([]float32, m.Dim)
		vectors = append(vectors, v)
		statuses = append(statuses, b.Status)
		rsIds = append(rsIds, b.RsId)
	}

	urnColumn := entity.NewColumnVarChar(UrnCol, urns)
	phaseColumn := entity.NewColumnVarChar(PhaseCol, phases)
	statusColumn := entity.NewColumnVarChar(StatusCol, statuses)
	reasonColumn := entity.NewColumnVarChar(ReasonCol, reasons)
	vectorColumn := entity.NewColumnFloatVector(VectorCol, m.Dim, vectors)
	rsIdColumn := entity.NewColumnInt32(RidCol, rsIds)

	e, err := m.Client.Upsert(m.Ctx, collectionName, "", urnColumn, phaseColumn, statusColumn, reasonColumn, rsIdColumn, vectorColumn)
	if err != nil {
		slog.Error("failed to upsert build data:", slog.Any("error:", err))
		return 0, err
	}

	if err := m.Client.Flush(m.Ctx, collectionName, false); err != nil {
		slog.Error("failed to flush data:", slog.Any("error:", err))
	}

	return e.Len(), nil
}

func (m *MilvusService) GetBuildLog(docId string, ids []string) ([]*BuildLog, error) {
	collectionName := fmt.Sprintf("%s%s", BuildLogPrefix, docId)

	err := m.LoadCollection(collectionName)
	if err != nil {
		if strings.Contains(err.Error(), "index not found") {
			return nil, nil
		}
		return nil, err
	}

	urnColumn := entity.NewColumnVarChar(UrnCol, ids)

	results, err := m.Client.Get(m.Ctx, collectionName, urnColumn)
	if err != nil {
		return nil, err
	}

	segments := []*BuildLog{}

	for _, field := range results {

		switch field.Name() {
		case UrnCol:
			c, ok := field.(*entity.ColumnVarChar)
			if !ok {
				continue
			}
			for i := 0; i < c.Len(); i++ {
				v, e := c.ValueByIdx(i)
				if e != nil {
					continue
				}

				if len(segments) < i+1 {
					segments = append(segments, &BuildLog{ChunkId: v})
				} else {
					segments[i].ChunkId = v
				}
			}
		case StatusCol:
			c, ok := field.(*entity.ColumnVarChar)
			if !ok {
				continue
			}
			for i := 0; i < c.Len(); i++ {
				v, e := c.ValueByIdx(i)
				if e != nil {
					continue
				}

				if len(segments) < i+1 {
					segments = append(segments, &BuildLog{Status: v})
				} else {
					segments[i].Status = v
				}
			}
		case ReasonCol:
			c, ok := field.(*entity.ColumnVarChar)
			if !ok {
				continue
			}
			for i := 0; i < c.Len(); i++ {
				v, e := c.ValueByIdx(i)
				if e != nil {
					continue
				}

				if len(segments) < i+1 {
					segments = append(segments, &BuildLog{Reason: v})
				} else {
					segments[i].Reason = v
				}
			}
		case RidCol:
			c, ok := field.(*entity.ColumnInt32)
			if !ok {
				continue
			}
			for i := 0; i < c.Len(); i++ {
				v, e := c.ValueByIdx(i)
				if e != nil {
					continue
				}

				if len(segments) < i+1 {
					segments = append(segments, &BuildLog{RsId: v})
				} else {
					segments[i].RsId = v
				}
			}
		default:
			slog.Warn("GetBuildLog ignore field:", slog.Any("field:", field.Name()))

		}
	}
	return segments, nil
}
